from core.data.data_handler import DataHandler
from core.memory.session_manager import SessionManager
from api.user_input_handler import UserInputHandler
from core.query.query_executor import QueryExecutor
from core.query.prompt.prompt_builder import PromptBuilder
from presentation.result_presenter import ResultPresenter
from core.utils.logging_handler import LoggingHandler
from core.query.query_refiner import QueryRefiner
from core.model.model_caller import ModelCaller
import pandas as pd

class QueryProcessor:
    def __init__(self, user_input, file_path, api_key, base_url, local_model_name, remote_model_name, explanation_file=None, use_remote=True):
        self.logger = LoggingHandler().get_logger()
        self.logger.info("初始化 QueryProcessor")
        self.data_handler = DataHandler(file_path, explanation_file)
        self.session_manager = SessionManager()
        self.user_input_handler = None
        self.query_executor = None
        self.result_presenter = None
        self.user_input = user_input
        self.api_key = api_key
        self.base_url = base_url
        self.local_model_name = local_model_name
        self.remote_model_name = remote_model_name
        self.explanation_file = explanation_file
        self.use_remote = use_remote
        self.result_summary = None
        self.result_txt = None

        # 初始化 ModelCaller
        self.model_caller = ModelCaller(
            api_key=self.api_key,
            base_url=self.base_url,
            local_model_name=self.local_model_name,
            remote_model_name=self.remote_model_name
        )

    def process(self):
        try:
            self.logger.info(f"处理用户输入: {self.user_input}")
            self.data_handler.load_data()
            self.logger.info("数据加载完成")
            df = self.data_handler.df
            
            # 初始化 QueryRefiner 并传递 ModelCaller
            refiner = QueryRefiner(
                model_caller=self.model_caller,
                df=df
            )

            # 调用模型精确化用户查询，决定使用本地还是远程模型
            refined_user_input = refiner.refine_query_with_model(self.user_input, use_remote=self.use_remote)
            self.logger.info(f"精确化后的用户输入: {refined_user_input}")
            
            # 判断是否为闲聊应答，如果是闲聊应答则直接返回结果
            if "闲聊应答" in refined_user_input:
                self.result_txt = refined_user_input
                print(refined_user_input)
                self.session_manager.add_to_history(self.user_input, None, None, refined_user_input, None)
                return
            
            # 处理用户输入，创建 UserInputHandler
            self.user_input_handler = UserInputHandler(
                user_input=refined_user_input, 
                api_key=self.api_key,
                base_url=self.base_url,
                local_model_name=self.local_model_name,
                remote_model_name=self.remote_model_name,
                explanations=self.data_handler.get_all_terms(), 
                conversation_history=self.session_manager.format_for_model(),
                use_remote=self.use_remote
            )



            query_code_or_response = self.user_input_handler.parse_query(df)

            generated_code = None
            execution_result = None
            natural_language_response = None
            summary = None
            
            if "result =" in query_code_or_response:
                self.query_executor = QueryExecutor(query_code_or_response)
                execution_result = self.query_executor.execute(df)
                generated_code = query_code_or_response
                self.result_presenter = ResultPresenter(execution_result)
                # self.result_presenter.display()
                # 新增功能：将查询结果和原始问题发送给大模型，生成总结
                summary = self.get_summary_from_model(execution_result, generated_code)
                print("Summary of the query result:")
                print(summary)
                self.result_summary = summary
            else:
                natural_language_response = query_code_or_response
                self.result_txt = query_code_or_response
                print(query_code_or_response)

            self.session_manager.add_to_history(self.user_input, generated_code, execution_result, natural_language_response, summary)
            
        except Exception as e:
            self.logger.error(f"处理过程中发生异常: {e}", exc_info=True)

    def get_summary_from_model(self, result, generated_code=None):
        """调用大模型生成对查询结果的总结"""
        result_as_str = result.to_string() if isinstance(result, pd.DataFrame) else str(result)
        
        # 创建用于生成总结的提示词
        summary_prompt_builder = PromptBuilder(
            df=None,  # 空，因为不需要列信息
            explanations=None,  # 不需要术语解释
            conversation_history=self.session_manager.format_for_model(),
            user_input=self.user_input
        )
        
        summary_prompt = summary_prompt_builder.build_summary_prompt(result_str=result_as_str, generated_code=generated_code)

        # 使用 UserInputHandler 调用模型生成总结
        summary_handler = UserInputHandler(
            user_input=summary_prompt,
            api_key=self.api_key,
            base_url=self.base_url,
            local_model_name=self.local_model_name,
            remote_model_name=self.remote_model_name,
            use_remote=self.use_remote
        )
        
        summary = summary_handler.parse_query(custom_prompt=summary_prompt)
        return summary
